#include <utility>
#include <fstream>
#include <algorithm>
#include "gc.h"

void debug_graph_connection (Graph3D const& graph, int startSliceNo, 
			     int layerNum);

bool isEqual (std::vector<int>& set0, std::vector<int>& set1)
{
  if (set0.size() != set1.size()) {
    return false;
  }
  sort(set0.begin(), set0.end());
  sort(set1.begin(), set1.end());
  std::vector<int>::const_iterator itr0 = set0.begin();
  std::vector<int>::const_iterator itr1 = set1.begin();
  while (itr0 != set0.end()) {
    if (*itr0 != *itr1) {
      return false;
    }
    itr0++;
    itr1++;
  }
  return true;
}



bool isContained (std::vector<int> const& set, int element) {
  for (std::vector<int>::const_iterator itr = set.begin(); 
       itr != set.end(); itr++) {
    if (*itr == element) {
      return true;
    }
  }
  return false;
}



std::vector<int> extractComponentNoSet (SegmentationGraph const& segGraph, 
					int layerNo, Graph3D const& graph)
{
  std::vector<int> ret;
  for (SegGraphSegNodeSet::const_iterator itr = segGraph.segNodeSet.begin(); 
       itr != segGraph.segNodeSet.end(); itr++) {
    if (itr->layerNo == layerNo) {
      ret.push_back(graph.nodes[itr->layerNo][itr->nodeNo].componentNo);
    }
  }
  return ret;
}



// set0 is better the bigger one (than set1)
bool isIntersected (std::vector<int> const& set0, std::vector<int> const& set1)
{
  for (std::vector<int>::const_iterator itr = set0.begin(); 
       itr != set0.end(); itr++) {
    if (isContained(set1, *itr)) {
      return true;
    }
  }
  return false;
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void initialization (Graph3D& graph, std::string const& initialFilePath, 
		     std::vector<std::vector<int> >& startComponentNoSets)
{
  startComponentNoSets.clear();
  std::ifstream fs(initialFilePath.c_str());
  if (fs.is_open()) {		// For watershed
    while (fs.good()) {
      int truthComponentNo;
      int waterComponentNoNum = -1;
      fs >> truthComponentNo >> waterComponentNoNum;
      if (waterComponentNoNum > 0) {
	std::vector<int> set;
	set.reserve(waterComponentNoNum);
	for (int waterComponentNoNo = 0; 
	     waterComponentNoNo < waterComponentNoNum; waterComponentNoNo++) {
	  int waterComponentNo = -1;
	  fs >> waterComponentNo;
	  int nodeNo = -1;
	  for (GraphLayerNodeSet::const_iterator itr = graph.nodes[0].begin(); 
	       itr != graph.nodes[0].end(); itr++) {
	    if (itr->componentNo == waterComponentNo) {
	      nodeNo = itr->no;
	      break;
	    }
	  }
	  if (nodeNo == -1) {
	    std::cout << "Weird!" << std::endl;
	    exit(1);
	  }
	  if (truthComponentNo > 0 
	      && graph.nodes[0][nodeNo].status == NODE_NORMAL) {
	    set.push_back(waterComponentNo);
	  }
	  else if (truthComponentNo == 0) {
	    graph.nodes[0][nodeNo].status = NODE_DELETED;
	  }
	}
	if (!set.empty()) {
	  startComponentNoSets.push_back(set);
	}
      }
    }
    fs.close();
  }
  else {			// For truth
    for (GraphLayerNodeSet::const_iterator nitr = graph.nodes[0].begin(); 
	 nitr != graph.nodes[0].end(); nitr++) {
      if (nitr->status == NODE_NORMAL && nitr->componentNo > 0) {
	std::vector<int> set(1, nitr->componentNo);
	startComponentNoSets.push_back(set);
      }
    }
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
// Return order # of desired segmentation
int greedy (Graph3D& graph, std::string const& initialFilePath, 
	    std::vector<int>& targetStartComponentNoSet, 
	    int startSliceNo, int startLayerNo, int layerNum, 
	    SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	    double& fwdMaxFlow, double& bwdMaxFlow, 
	    VerboseLevel verbose)
{
  int order = 0;
  std::vector<std::vector<int> > startComponentNoSets; 
  std::map<std::vector<int>, double> fwdMaxFlowTable;
  std::map<std::vector<int>, double> bwdMaxFlowTable;
  std::map<std::vector<int>, SegmentationGraph> fwdSegGraphTable;
  std::map<std::vector<int>, SegmentationGraph> bwdSegGraphTable;
  while (true) {
    startComponentNoSets.clear();
    fwdMaxFlowTable.clear();
    bwdMaxFlowTable.clear();
    fwdSegGraphTable.clear();
    bwdSegGraphTable.clear();
    initialization(graph, initialFilePath, startComponentNoSets);
    if (startComponentNoSets.size() <= 0) {
      if (verbose != VERBOSE_NONE) {
	std::cout << "Fail to track target starting components..." << std::endl;
      }
      exit(1);
    }
    double maxMaxFlow = -1.0;
    std::vector<int> maxComponentNoSet;
    for (std::vector<std::vector<int> >::const_iterator itr = 
	   startComponentNoSets.begin(); 
	 itr != startComponentNoSets.end(); itr++) {
      double curFwdMaxFlow, curBwdMaxFlow;
      SegmentationGraph curFwdSegGraph, curBwdSegGraph;
      segment(graph, *itr, startLayerNo, layerNum, 
	      curFwdSegGraph, curBwdSegGraph, curFwdMaxFlow, curBwdMaxFlow, 
	      verbose);
      fwdMaxFlowTable[*itr] = curFwdMaxFlow;
      bwdMaxFlowTable[*itr] = curBwdMaxFlow;
      fwdSegGraphTable[*itr] = curFwdSegGraph;
      bwdSegGraphTable[*itr] = curBwdSegGraph;
      double curMaxFlow = curFwdMaxFlow + curBwdMaxFlow;
      if (maxMaxFlow < curMaxFlow) {
	maxMaxFlow = curMaxFlow;
	maxComponentNoSet = *itr;
      }
    }
    if (isEqual(maxComponentNoSet, targetStartComponentNoSet)) {
      // // For debug
      // debug_graph_connection(graph, startSliceNo, layerNum);
      // exit(1);
      // // ~ For debug
      fwdMaxFlow = fwdMaxFlowTable[targetStartComponentNoSet];
      bwdMaxFlow = bwdMaxFlowTable[targetStartComponentNoSet];
      fwdSegGraph = fwdSegGraphTable[targetStartComponentNoSet];
      bwdSegGraph = bwdSegGraphTable[targetStartComponentNoSet];
      return order;
    }
    else {
      SegmentationGraph maxSegGraph = 
	fwdSegGraphTable[maxComponentNoSet].getUnion(bwdSegGraphTable[maxComponentNoSet]);
      graph.changeNodeStatus(maxSegGraph, NODE_DELETED);
      if (verbose == VERBOSE_INFO) {
	std::cout << "Order " << order << ": removed starting component ";
	for (std::vector<int>::const_iterator itr = maxComponentNoSet.begin(); 
	     itr != maxComponentNoSet.end(); itr++) {
	  std::cout << *itr << " ";
	}
	std::cout << "(f = " 
		  << fwdMaxFlowTable[maxComponentNoSet] + bwdMaxFlowTable[maxComponentNoSet] 
		  << ")" << std::endl;
	for (SegGraphSegNodeSet::const_iterator itr = 
	       maxSegGraph.segNodeSet.begin(); 
	     itr != maxSegGraph.segNodeSet.end(); itr++) {
	  std::cout << "remove: layer " << itr->layerNo + startSliceNo 
		    << ", component " 
		    << graph.nodes[itr->layerNo][itr->nodeNo].componentNo 
		    << std::endl;
	}
      }
    }
    order++;
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
// Output maxflow segmentations in a row
// No targets
// Return total # of maxflow segmentations
int greedy (Graph3D& graph, std::string const& initialFilePath, 
	    int startSliceNo, int startLayerNo, int layerNum, 
	    VerboseLevel verbose)
{
  int order = 0;
  std::vector<std::vector<int> > startComponentNoSets;
  std::map<std::vector<int>, double> fwdMaxFlowTable;
  std::map<std::vector<int>, double> bwdMaxFlowTable;
  std::map<std::vector<int>, SegmentationGraph> fwdSegGraphTable;
  std::map<std::vector<int>, SegmentationGraph> bwdSegGraphTable;
  while (true) {
    startComponentNoSets.clear();
    fwdMaxFlowTable.clear();
    bwdMaxFlowTable.clear();
    fwdSegGraphTable.clear();
    bwdSegGraphTable.clear();    
    initialization(graph, initialFilePath, startComponentNoSets);
    if (startComponentNoSets.size() == 0) {
      return order;
    }
    double maxMaxFlow = -1.0;
    std::vector<int> maxComponentNoSet;
    for (std::vector<std::vector<int> >::const_iterator itr = 
	   startComponentNoSets.begin(); 
	 itr != startComponentNoSets.end(); itr++) {
      double curFwdMaxFlow, curBwdMaxFlow;
      SegmentationGraph curFwdSegGraph, curBwdSegGraph;
      segment(graph, *itr, startLayerNo, layerNum, 
	      curFwdSegGraph, curBwdSegGraph, curFwdMaxFlow, curBwdMaxFlow, 
	      verbose);
      fwdMaxFlowTable[*itr] = curFwdMaxFlow;
      bwdMaxFlowTable[*itr] = curBwdMaxFlow;
      fwdSegGraphTable[*itr] = curFwdSegGraph;
      bwdSegGraphTable[*itr] = curBwdSegGraph;
      double curMaxFlow = curFwdMaxFlow + curBwdMaxFlow;
      if (maxMaxFlow < curMaxFlow) {
	maxMaxFlow = curMaxFlow;
	maxComponentNoSet = *itr;
      }
    }
    SegmentationGraph maxSegGraph = 
      fwdSegGraphTable[maxComponentNoSet].getUnion(bwdSegGraphTable[maxComponentNoSet]);
    graph.changeNodeStatus(maxSegGraph, NODE_DELETED);
    if (verbose == VERBOSE_INFO) {
      std::cout << "Order " << order << ": removed starting component ";
      for (std::vector<int>::const_iterator itr = maxComponentNoSet.begin(); 
	   itr != maxComponentNoSet.end(); itr++) {
	std::cout << *itr << " ";
      }
      std::cout << "(f = " 
		<< fwdMaxFlowTable[maxComponentNoSet] + bwdMaxFlowTable[maxComponentNoSet] 
		<< ")" << std::endl;
      for (SegGraphSegNodeSet::const_iterator itr = 
	     maxSegGraph.segNodeSet.begin(); 
	   itr != maxSegGraph.segNodeSet.end(); itr++) {
	std::cout << "remove: layer " << itr->layerNo + startSliceNo 
		  << ", component " 
		  << graph.nodes[itr->layerNo][itr->nodeNo].componentNo 
		  << std::endl;
      }
    }
    order++;
  }
}
